from transformers import pipeline
import torch
from tqdm import tqdm

test_path = "/home/mbk/rubbish/sen_cls/data/ASAP_SENT/test.tsv"
model_path = "../roberta-base-finetuned-dianping-chinese"


classifier = pipeline("text-classification", model=model_path, truncation=True, max_length=512)

def predict_sentiment(text):
    try:
        result = classifier(text)
        return int(result[0]['score']/0.2)
    except Exception as e:
        print(text)
        print(e)


res = "index	prediction\n"

with open(test_path, 'r') as f:
    data = f.read().split('\n')[1:]
    for line in tqdm(data):
        if len(line) == 0:
            continue
        qid, text = tuple(line.split('\t'))
        label = predict_sentiment(text)
        res += f"{qid}\t{label+1}\n"

with open('ASAP_SENT.tsv', 'w') as f:
    f.write(res)
